from super_gradients import Trainer, setup_device
from super_gradients.common.object_names import Models
from super_gradients.training import dataloaders
from super_gradients.training.dataloaders.dataloaders import coco_detection_yolo_format_train, \
    coco_detection_yolo_format_val
from super_gradients.training import models
from super_gradients.training.losses import PPYoloELoss
from super_gradients.training.metrics import DetectionMetrics_050
from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback
import os

# os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

CHECKPOINT_DIR = '.\\test'


def train():
    dataset_params = {
        'data_dir': r'D:\mypro\yolo-nas\datasets\breast\yolo-darknet',
        'train_img_dir': 'train2007',
        'train_labels_dir': 'train2007',
        'val_img_dir': 'val2007',
        'val_labels_dir': 'val2007',
        'test_img_dir': 'test2007',
        'test_labels_dir': 'test2007',
        'classes': ['breast'],
        'cache_dir': r'D:\mypro\yolo-nas\datasets\breast\yolo-darknet',
        'input_dim': [640, 640]
    }

    train_data = coco_detection_yolo_format_train(
        dataset_params={
            'data_dir': dataset_params['data_dir'],
            'images_dir': dataset_params['train_img_dir'],
            'labels_dir': dataset_params['train_labels_dir'],
            'classes': dataset_params['classes'],
            'cache': True,
            'cache_dir': dataset_params['cache_dir'],
            'input_dim': dataset_params['input_dim']
        },
        dataloader_params={
            'batch_size': 8,
            'num_workers': 16
        }
    )

    val_data = coco_detection_yolo_format_val(
        dataset_params={
            'data_dir': dataset_params['data_dir'],
            'images_dir': dataset_params['val_img_dir'],
            'labels_dir': dataset_params['val_labels_dir'],
            'classes': dataset_params['classes'],
            'cache': True,
            'cache_dir': dataset_params['cache_dir'],
            'input_dim': dataset_params['input_dim']
        },
        dataloader_params={
            'batch_size': 8,
            'num_workers': 16
        }
    )

    model = models.get(Models.YOLO_NAS_S, num_classes=len(dataset_params['classes']), pretrained_weights="coco")

    train_params = {
        "average_best_models": True,
        "warmup_mode": "linear_epoch_step",
        "warmup_initial_lr": 1e-6,
        "lr_warmup_epochs": 3,
        "initial_lr": 2e-4,
        "lr_mode": "cosine",
        "cosine_final_lr_ratio": 0.1,
        "optimizer": "Adam",
        "optimizer_params": {"weight_decay": 0.0001},
        "zero_weight_decay_on_bias_and_bn": True,
        "ema": True,
        "ema_params": {"decay": 0.9, "decay_type": "threshold"},
        # ONLY TRAINING FOR 10 EPOCHS FOR THIS EXAMPLE NOTEBOOK
        "max_epochs": 300,
        "save_ckpt_epoch_list": range(10, 310, 10),
        "mixed_precision": True,
        "loss": PPYoloELoss(
            use_static_assigner=False,
            # NOTE: num_classes needs to be defined here
            num_classes=len(dataset_params['classes']),
            reg_max=16
        ),
        "valid_metrics_list": [
            DetectionMetrics_050(
                score_thres=0.1,
                top_k_predictions=300,
                # NOTE: num_classes needs to be defined here
                num_cls=len(dataset_params['classes']),
                normalize_targets=True,
                calc_best_score_thresholds=True,
                post_prediction_callback=PPYoloEPostPredictionCallback(
                    score_threshold=0.01,
                    nms_top_k=1000,
                    max_predictions=300,
                    nms_threshold=0.7
                )
            )
        ],
        "metric_to_watch": 'F1@0.50',
        "greater_metric_to_watch_is_better": True
    }
    setup_device(num_gpus=-1)
    trainer = Trainer(experiment_name='breast_detection', ckpt_root_dir=CHECKPOINT_DIR)
    trainer.train(model=model,
                  training_params=train_params,
                  train_loader=train_data,
                  valid_loader=val_data)


if __name__ == "__main__":
    train()
